{-# OPTIONS_HADDOCK prune, ignore-exports #-}
module TIChecker
(
  generateTyVarNames,

  hypotheses,

  inferType
)
where

import qualified Data.Map as Map

import TIParser  
import TITypes   

generateTyVarNames :: [String] 
generateTyVarNames = [a : if n == 0 then "" else show n | n <- [0..],
                                                          a <- ['a'..'z']]

hypotheses :: Expr -> [String] -> Env -> (Env, [String])  
hypotheses e@(Var n)   (tv:tvs) env = ((e, TyVar tv) : env, tvs)
hypotheses e@(App e1 e2) tvs    env = let (env1, tvs1)         = hypotheses e1 tvs env
                                          (env2, (t1:t2:tvs2)) = hypotheses e2 tvs1 env1
                                      in  ((e, TyVar t2) : (e1, Arrow (TyVar t1) (TyVar t2)) : (e2, TyVar t1) : env2, tvs2)
hypotheses e@(Abs n e1)  tvs    env = let ((x:env1), (tv:tvs')) = hypotheses e1 tvs env
                                      in  ((e, Arrow (TyVar tv) $ snd x) : (Var n, TyVar tv) : (x:env1), tvs')


inferType :: Expr -> [String] -> Environment -> Either String (Environment, [String]) 
inferType e@(Var n) t@(tv:tvs) env = case Map.lookup e env of
                                       Nothing -> return (Map.insert e (TyVar tv) env, tvs)
                                       Just _  -> return (env, t)
inferType e@(App e1 e2) t@(tv:tvs) env = do (env', tvs') <- inferType e2 tvs env
                                            let Just tp' = Map.lookup e2 env'
                                            case e1 of
                                              Var{} -> case Map.lookup e1 env' of
                                                         Nothing -> do return (Map.insert e (TyVar tv) (Map.insert e1 (Arrow tp' (TyVar tv)) env'), tvs')
                                                         Just tp -> do env'' <- unifyTypes (Arrow tp' (TyVar tv), tp) env'
                                                                       let Just (Arrow tp'' tp''') = Map.lookup e1 env''
                                                                       return (Map.insert e tp''' env'', tvs')
                                              App{} -> do (env'', tvs'') <- inferType e1 tvs' env'
                                                          let Just tp = Map.lookup e1 env''
                                                          env''' <- unifyTypes (tp, Arrow tp' (TyVar tv)) env''
                                                          let Just (Arrow tp'' tp''') = Map.lookup e1 env'''
                                                          return (Map.insert e tp''' env''', tvs'')
                                              Abs{} -> do (env'', tvs'') <- inferType e1 tvs' env'
                                                          let Just (Arrow tp1 tp2) = Map.lookup e1 env''
                                                          if areTypesCompatible tp1 tp'
                                                            then do env''' <- unifyTypes (tp1, tp') env''
                                                                    let Just (Arrow tp1' tp2') = Map.lookup e1 env'''
                                                                    return (Map.insert e tp2' env''', tvs'')
                                                            else fail ("ERROR: Can't apply \"" ++ show e1 ++ "\" to \"" ++ show e2 ++ "\". Incopatible types.")
inferType e@(Abs n e1) (tv:tvs) env = case Map.lookup (Var n) env of
                                        Nothing -> do (env', tvs') <- inferType e1 tvs env
                                                      let Just tp' = Map.lookup e1 env'
                                                      case Map.lookup (Var n) env' of
                                                        Nothing -> do let env'' = Map.insert (Var n) (TyVar tv) env'
                                                                      return (Map.insert e (Arrow (TyVar tv) tp') env'', tvs')
                                                        Just tp -> return (Map.insert e (Arrow tp tp') env', tvs')
                                        Just tp -> fail ("ERROR: Duplicate bound variable \"" ++ n ++ "\" in lambda-abstraction.")



unifyTypes :: (Type, Type) -> Environment -> Either String Environment 
unifyTypes (Const, Const) env               = return env
unifyTypes (TyVar n, t) env                 = return $ Map.map (substituteTyVar n t) env
unifyTypes (Arrow t1 t2, Arrow t1' t2') env = do env' <- unifyTypes (t1, t1') env
                                                 unifyTypes (t2, t2') env'
unifyTypes (t1, t2) _ = fail ("ERROR: Can't unify type (" ++ show t1 ++ ") with (" ++ show t2 ++ ").")


areTypesCompatible :: Type -> Type -> Bool  
areTypesCompatible Const          _              = True
areTypesCompatible TyVar{}        _              = True
areTypesCompatible (Arrow t1 t2) (Arrow t1' t2') = areTypesCompatible t1 t1' && areTypesCompatible t2 t2'
areTypesCompatible _             _               = False

substituteTyVar :: String -> Type -> Type -> Type   
substituteTyVar _ _ Const                  = Const
substituteTyVar n t (TyVar n') | n' == n   = t
                               | otherwise = TyVar n'
substituteTyVar n t (Arrow t1 t2)          = Arrow (substituteTyVar n t t1) (substituteTyVar n t t2)

